package org.example.batch.service;

import org.example.batch.dao.BatchMainTaskDao;
import org.example.batch.dao.BatchSubTaskDao;
import org.example.batch.dao.BatchTerResultDao;
import org.example.batch.exception.TableNotExistException;
import org.example.batch.exception.TableNotReadyException;
import org.example.batch.model.BatchSubTaskDto;
import org.example.batch.pojo.BatchMainTable;
import org.example.batch.pojo.BatchSubTable;
import org.example.batch.service.threadPool.CustomThreadPool;
import org.example.batch.task.BatchRunRecordTask;
import org.example.batch.util.constant.BatchTerSubStatus;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.ApplicationContext;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.CountDownLatch;

/**
 * 子任务逻辑处理类
 *
 * @author hero生仔
 * @date 2025/7/26
 */
@Service
public class BatchSubTaskService {

    private static final Logger LOGGER = LoggerFactory.getLogger(BatchSubTaskService.class);

    @Autowired
    private BatchSubTaskDao batchSubTaskDao;

    @Autowired
    private BatchMainTaskDao batchMainTaskDao;

    @Autowired
    private BatchTerResultDao batchTerResultDao;

    @Autowired
    private ApplicationContext applicationContext; // 注入Spring上下文

    private BatchMainTable batchMainTable;

    private BatchSubTable batchSubTable;

    private List<BatchSubTable> batchSubTableList = new ArrayList<>();

    /**
     * 可用线程数
     */
    @Value("${ter.threadCount}")
    private int threadCount;

    /**
     * 处理子任务
     *
     * @param batchSubTaskDto 子任务
     */
    public void handleBatchSubTask(BatchSubTaskDto batchSubTaskDto) throws InterruptedException {
        // 校验子任务，主任务是否存在且子任务是否可以启动
        checkMainAndSubTable(batchSubTaskDto);
        // 设置为 RUNNING 状态
        batchSubTable.setStatus(BatchTerSubStatus.RUNNING.name());

        // 初始化子任务线程信息
        if (batchSubTableList.isEmpty()) {
            initBatchSubTableThreadInfo(batchSubTable);
        }
        // 子线程设置为 RUNNING 状态
        for (BatchSubTable subTable : batchSubTableList) {
            subTable.setStatus(BatchTerSubStatus.RUNNING.name());
            subTable.setErrorMsg(null);
        }

        // 执行线程
        CountDownLatch countDownLatch = new CountDownLatch(threadCount);
        for (BatchSubTable subTableThread : batchSubTableList) {
            BatchRunRecordTask batchRunRecordTask = applicationContext.getBean(BatchRunRecordTask.class);
            batchRunRecordTask.setBatchMainTable(batchMainTable);
            batchRunRecordTask.setBatchSubTableThread(subTableThread);
            batchRunRecordTask.setCountDownLatch(countDownLatch);
            CustomThreadPool.execute(batchRunRecordTask);
        }
        // 等待线程全部执行完毕
        countDownLatch.await();

        LOGGER.info("聚合子线程到子任务信息");
        // 子线程的数据聚合到子任务
        aggregateBatchSubTable();
    }

    /**
     * 子线程的数据聚合到子任务
     */
    private void aggregateBatchSubTable() {
        List<BatchSubTable> batchSubTableThreadList = batchSubTaskDao.queryBatchSubTableInfos(batchSubTable.getSubTaskId());
        int runningCount = 0;
        for (int i = 0; i < batchSubTableThreadList.size(); i++) {
            String status = batchSubTableThreadList.get(i).getStatus();
            runningCount += batchSubTableThreadList.get(i).getAlreadyRunningCount();
            // 如果其中一个子线程不成功，那么该子任务就不成功
            if (!BatchTerSubStatus.SUCCESS.name().equalsIgnoreCase(status)) {
                batchSubTable.setStatus(status);
            }
            // 遍历最后一个线程的时候，如果已运行数和输入表总数相当：代表子任务成功
            if (i == batchSubTableThreadList.size() - 1) {
                batchSubTable.setAlreadyRunningCount(runningCount);
                if (runningCount == batchSubTable.getInputTableTotalCount()) {
                    batchSubTable.setLastQueryKey(batchSubTableThreadList.get(i).getLastQueryKey());
                    batchSubTable.setStatus(BatchTerSubStatus.SUCCESS.name());
                }
            }
        }
        // 更新子任务状态和已运行数量
        batchSubTaskDao.updateSubTableMsg(batchSubTable);
    }

    /**
     * 初始化子任务的线程信息
     *
     * @param batchSubTable 子任务
     */
    private void initBatchSubTableThreadInfo(BatchSubTable batchSubTable) {
        int totalCount = batchSubTable.getInputTableTotalCount();
        // 每个线程处理的记录总数
        int perThreadHandleCount = totalCount / threadCount;
        // 余数：交给最后一个线程处理
        int remainCount = totalCount % threadCount;
        int start = 1;
        // 子线程也是子任务：子线程数据需要进行预持久化
        for (int i = 0; i < threadCount; i++) {
            int end = start + perThreadHandleCount - 1;
            if (i == threadCount - 1) {
                end += remainCount;
                perThreadHandleCount = perThreadHandleCount + remainCount;
            }
            BatchSubTable batchSubTableThread = new BatchSubTable(batchSubTable.getSubTaskId() + "-t" + (i + 1));
            batchSubTableThread.addBaseInfo(batchSubTable, start, end, perThreadHandleCount, start);
            batchSubTableList.add(batchSubTableThread);

            // 更新 起始索引
            start = end + 1;
        }
        batchSubTaskDao.batchInsertSubTableInfos(batchSubTableList);
    }

    /**
     * 校验主/子任务
     *
     * @param batchSubTaskDto 子任务DTO
     */
    private void checkMainAndSubTable(BatchSubTaskDto batchSubTaskDto) {
        String subTaskId = batchSubTaskDto.getSubTaskId();
        batchSubTable = batchSubTaskDao.queryBatchSubTableInfo(subTaskId);
        if (null == batchSubTable) {
            throw new TableNotExistException("子任务不存在");
        }
        String mainTaskId = batchSubTable.getMainTaskId();
        batchMainTable = batchMainTaskDao.queryBatchMainTableInfo(mainTaskId);
        if (null == batchMainTable) {
            throw new TableNotExistException("主任务不存在");
        }
        // 子任务需要为INIT或者CALL_FAIL状态才能启动
        if (!Arrays.asList(BatchTerSubStatus.INIT.name(), BatchTerSubStatus.CALL_FAIL.name()).contains(batchSubTable.getStatus())) {
            throw new TableNotReadyException("子任务不处于INIT状态或者CALL_FAIL状态");
        }
        // 查询子任务是否有子线程
        batchSubTableList = batchSubTaskDao.queryBatchSubTableInfos(batchSubTable.getSubTaskId());
    }
}
